import tensorflow as tf
from numpy import random
from tensorflow.python.keras.backend import set_session
from model.train_model import model_fn
from keras.applications.resnet50 import preprocess_input
from PIL import Image
import requests as req
from io import BytesIO
import numpy as np
import os
import random
import re
import json

# 全局配置文件
tf.compat.v1.app.flags.DEFINE_string('test_data_local', './test_image', '测试图片文件夹')
tf.compat.v1.app.flags.DEFINE_string('data_local', './jupyterLab/data/base_data', '训练图片文件夹')
tf.compat.v1.app.flags.DEFINE_integer('num_classes', 40, '垃圾分类数目')
tf.compat.v1.app.flags.DEFINE_integer('input_size', 224, '模型输入图片大小')
tf.compat.v1.app.flags.DEFINE_integer('batch_size', 16, '图片批处理大小')
tf.compat.v1.app.flags.DEFINE_float('learning_rate', 1e-4, '学习率')
tf.compat.v1.app.flags.DEFINE_integer('max_epochs', 4, '步长')
tf.compat.v1.app.flags.DEFINE_string('train_local', './output_model/', '训练输出文件夹')
tf.compat.v1.app.flags.DEFINE_integer('keep_weights_file_num', 20, '如果设置为-1，则文件保持的最大权重数表示无穷大')
FLAGS = tf.compat.v1.app.flags.FLAGS


# import serial
# OPTIONAL: control usage of GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = 0.7
sess = tf.compat.v1.Session(config=config)

h5_weights_path = '../jupyterLab/output_model/best.h5'


def init_artificial_neural_network(sess):
    set_session(sess)
    model = model_fn(FLAGS)
    model.load_weights(h5_weights_path, by_name=True)
    return model


## 测试图片
def prediction_result_from_img(model,imgurl):
    # 加载分类数据
    try:
        with open("../jupyterLab/data/garbage_classify_rule.json", 'r') as load_f:
            load_dict = json.load(load_f)
        if re.match(r'^https?:/{2}\w.+$', imgurl):
            test_data = preprocess_img_from_Url(imgurl, FLAGS.input_size)
        else:
            test_data = preprocess_img(imgurl, FLAGS.input_size)

        if test_data != 0:
            tta_num = 5
            predictions = [0 * tta_num]
            for i in range(tta_num):
                x_test = test_data[i]
                x_test = x_test[np.newaxis, :, :, :]
                prediction = model.predict(x_test)[0]
                # print(prediction)
                predictions += prediction
            pred_label = np.argmax(predictions, axis=0)
            print('-------深度学习垃圾分类预测结果----------')
            print(pred_label)
            print(load_dict[str(pred_label)])
            print('-------深度学习垃圾分类预测结果--------')
            return load_dict[str(pred_label)]
        else:
            print('-------文件读取错误----------')
            return False
    except Exception as e:
        print('发生了异常-prediction：', e)


# 本地路径获取图片信息
def preprocess_img(img_path,img_size):
    try:
        img = Image.open(img_path)
        # if img.format:
        # resize_scale = img_size / max(img.size[:2])
        # img = img.resize((int(img.size[0] * resize_scale), int(img.size[1] * resize_scale)))
        img = img.resize((256, 256))
        img = img.convert('RGB')
        # img.show()
        img = np.array(img)
        imgs = []
        for _ in range(10):
            i = random.randint(0, 32)
            j = random.randint(0, 32)
            imgg = img[i:i + 224, j:j + 224]
            imgg = preprocess_input(imgg)
            imgs.append(imgg)
        return imgs
    except Exception as e:
        print('发生了异常data_process：', e)
        return 0




# url获取图片数组信息
def preprocess_img_from_Url(img_path,img_size):
    try:
        response = req.get(img_path)
        img = Image.open(BytesIO(response.content))
        img = img.resize((256, 256))
        img = img.convert('RGB')
        # img.show()
        img = np.array(img)
        imgs = []
        for _ in range(10):
            i = random.randint(0, 32)
            j = random.randint(0, 32)
            imgg = img[i:i + 224, j:j + 224]
            imgg = preprocess_input(imgg)
            imgs.append(imgg)
        return imgs
    except Exception as e:
        print('发生了异常data_process：', e)
        return 0